import numpy as np
import math
import matplotlib.pyplot as plt
# import theano

import A2C.parameters as parameters


class Env:
    def __init__(self, pa, nw_len_seqs=None, nw_size_seqs=None,
                 seed=42, render=False, repre='image', end='no_new_resource'):

        self.pa = pa
        self.render = render
        self.repre = repre  # image or compact representation
        self.end = end  # termination type, 'no_new_job' or 'all_done'

        self.nw_dist = pa.dist.bi_model_dist

        self.curr_time = 0

        self.weight = 0.5

        # set up random seed
        if self.pa.unseen:
            np.random.seed(314159)
        else:
            np.random.seed(seed)

        if nw_len_seqs is None or nw_size_seqs is None:
            # generate new work
            self.nw_len_seqs, self.nw_size_seqs = \
                self.generate_sequence_work(self.pa.simu_len * self.pa.num_ex)

            self.workload = np.zeros(pa.num_res)
            for i in range(pa.num_res):
                self.workload[i] = \
                    np.sum(self.nw_size_seqs[:, i] * self.nw_len_seqs) / \
                    float(pa.res_slot) / \
                    float(len(self.nw_len_seqs))
                print("Load on # " + str(i) + " resource dimension is " + str(self.workload[i]))
            self.nw_len_seqs = np.reshape(self.nw_len_seqs,
                                          [self.pa.num_ex, self.pa.simu_len])
            self.nw_size_seqs = np.reshape(self.nw_size_seqs,
                                           [self.pa.num_ex, self.pa.simu_len, self.pa.num_res])
        else:
            self.nw_len_seqs = nw_len_seqs
            self.nw_size_seqs = nw_size_seqs

        self.seq_no = 0  # which example sequence
        self.seq_idx = 0  # index in that sequence

        # initialize system
        self.resource = Resource(pa)
        self.resource_slot = DisasterResourceSlot(pa)
        self.resource_backlog = DisasterResourceBacklog(pa)
        self.resource_record = DisasterResourceRecord()
        self.extra_info = ExtraInfo(pa)

    def generate_sequence_work(self, simu_len):

        nw_len_seq = np.zeros(simu_len, dtype=int)
        nw_size_seq = np.zeros((simu_len, self.pa.num_res), dtype=int)

        for i in range(simu_len):

            if np.random.rand() < self.pa.new_job_rate:  # a new resource comes

                nw_len_seq[i], nw_size_seq[i, :] = self.nw_dist()

        return nw_len_seq, nw_size_seq

    def get_new_resource_from_seq(self, seq_no, seq_idx):
        temp = np.random.rand(0,2)
        if temp > 1:
            weight = 0.4
        else:
            weight = 0.6

        new_resource = DisasterResource(res_vec=self.nw_size_seqs[seq_no, seq_idx, :],
                      resource_len=self.nw_len_seqs[seq_no, seq_idx],
                      resource_id=len(self.resource_record.record),
                      enter_time=self.curr_time,weight=weight)
        return new_resource

    def observe(self):
        if self.repre == 'image':

            backlog_width = int(math.ceil(self.pa.backlog_size / float(self.pa.time_horizon)))

            image_repr = np.zeros((self.pa.network_input_height, self.pa.network_input_width))

            ir_pt = 0

            for i in range(self.pa.num_res):

                image_repr[:, ir_pt: ir_pt + self.pa.res_slot] = self.resource.canvas[i, :, :]
                ir_pt += self.pa.res_slot

                for j in range(self.pa.num_nw):

                    if self.resource_slot.slot[j] is not None:  # fill in a block of work
                        image_repr[: self.resource_slot.slot[j].len, ir_pt: ir_pt + self.resource_slot.slot[j].res_vec[i]] = 1

                    ir_pt += self.pa.max_job_size

            image_repr[: int(self.resource_backlog.curr_size / backlog_width),
            ir_pt: ir_pt + backlog_width] = 1
            if self.resource_backlog.curr_size % backlog_width > 0:
                image_repr[int(self.resource_backlog.curr_size / backlog_width),
                ir_pt: ir_pt + self.resource_backlog.curr_size % backlog_width] = 1
            ir_pt += backlog_width

            image_repr[:, ir_pt: ir_pt + 1] = self.extra_info.time_since_last_new_resource / \
                                              float(self.extra_info.max_tracking_time_since_last_resource)
            ir_pt += 1

            assert ir_pt == image_repr.shape[1]

            return image_repr.ravel()[np.newaxis, :]
            # return image_repr

    def plot_state(self):
        plt.figure("screen", figsize=(20, 5))

        skip_row = 0

        for i in range(self.pa.num_res):

            plt.subplot(self.pa.num_res,
                        1 + self.pa.num_nw + 1,  # first +1 for current work, last +1 for backlog queue
                        i * (self.pa.num_nw + 1) + skip_row + 1)  # plot the backlog at the end, +1 to avoid 0

            plt.imshow(self.resource.canvas[i, :, :], interpolation='nearest', vmax=1)

            for j in range(self.pa.num_nw):

                resource_slot = np.zeros((self.pa.time_horizon, self.pa.max_job_size))
                if self.resource_slot.slot[j] is not None:  # fill in a block of work
                    resource_slot[: self.resource_slot.slot[j].len, :self.resource_slot.slot[j].res_vec[i]] = 1

                plt.subplot(self.pa.num_res,
                            1 + self.pa.num_nw + 1,  # first +1 for current work, last +1 for backlog queue
                            1 + i * (
                                        self.pa.num_nw + 1) + j + skip_row + 1)  # plot the backlog at the end, +1 to avoid 0

                plt.imshow(resource_slot, interpolation='nearest', vmax=1)

                if j == self.pa.num_nw - 1:
                    skip_row += 1

        skip_row -= 1
        backlog_width = int(math.ceil(self.pa.backlog_size / float(self.pa.time_horizon)))
        backlog = np.zeros((self.pa.time_horizon, backlog_width))

        backlog[: self.resource_backlog.curr_size / backlog_width, : backlog_width] = 1
        backlog[self.resource_backlog.curr_size / backlog_width, : self.resource_backlog.curr_size % backlog_width] = 1

        plt.subplot(self.pa.num_res,
                    1 + self.pa.num_nw + 1,  # first +1 for current work, last +1 for backlog queue
                    self.pa.num_nw + 1 + 1)

        plt.imshow(backlog, interpolation='nearest', vmax=1)

        plt.subplot(self.pa.num_res,
                    1 + self.pa.num_nw + 1,  # first +1 for current work, last +1 for backlog queue
                    self.pa.num_res * (self.pa.num_nw + 1) + skip_row + 1)  # plot the backlog at the end, +1 to avoid 0

        extra_info = np.ones((self.pa.time_horizon, 1)) * \
                     self.extra_info.time_since_last_new_resource / \
                     float(self.extra_info.max_tracking_time_since_last_resource)

        plt.imshow(extra_info, interpolation='nearest', vmax=1)

        plt.show()  # manual
        # plt.pause(0.01)  # automatic


    def get_reward(self):

        reward = 0
        for resource in self.resource.running_resource:
            reward += self.pa.delay_penalty / float(resource.len) * resource.weight

        for resource in self.resource_slot.slot:
            if resource is not None:
                reward += self.pa.hold_penalty / float(resource.len) * resource.weight

        for resource in self.resource_backlog.backlog:
            if resource is not None:
                reward += self.pa.dismiss_penalty / float(resource.len) * resource.weight

        return reward

    def step(self, a, repeat=False):

        status = None

        done = False
        reward = 0
        info = None

        if a == self.pa.num_nw:  # explicit void action
            status = 'MoveOn'
        elif self.resource_slot.slot[a] is None:  # implicit void action
            # if self.seq_idx >= self.pa.simu_len and \
            # len(self.machine.running_job) > 0 and \
            # all(s is None for s in self.job_backlog.backlog):
            # ob, reward, done, info = self.step(a + 1, repeat=True)
            # return ob, reward, done, info
            # else:
            status = 'MoveOn'
        else:
            allocated = self.resource.allocate_resource(self.resource_slot.slot[a], self.curr_time)
            if not allocated:  # implicit void action
                status = 'MoveOn'
            else:
                status = 'Allocate'

        if status == 'MoveOn':
            self.curr_time += 1
            self.resource.time_proceed(self.curr_time)
            self.extra_info.time_proceed()

            if self.end == "no_new_resource":  # end of new job sequence
                if self.seq_idx >= self.pa.simu_len:
                    done = True
            elif self.end == "all_done":  # everything has to be finished
                if self.seq_idx >= self.pa.simu_len and \
                        len(self.resource.running_resource) == 0 and \
                        all(s is None for s in self.resource_slot.slot) and \
                        all(s is None for s in self.resource_backlog.backlog):
                    done = True
                elif self.curr_time > self.pa.episode_max_length:  # run too long, force termination
                    done = True

            if not done:

                if self.seq_idx < self.pa.simu_len:  # otherwise, end of new job sequence, i.e. no new jobs
                    new_resource = self.get_new_resource_from_seq(self.seq_no, self.seq_idx)

                    if new_resource.len > 0:  # a new job comes

                        to_backlog = True
                        for i in range(self.pa.num_nw):
                            if self.resource_slot.slot[i] is None:  # put in new visible job slots
                                self.resource_slot.slot[i] = new_resource
                                self.resource_record.record[new_resource.id] = new_resource
                                to_backlog = False
                                break

                        if to_backlog:
                            if self.resource_backlog.curr_size < self.pa.backlog_size:
                                self.resource_backlog.backlog[self.resource_backlog.curr_size] = new_resource
                                self.resource_backlog.curr_size += 1
                                self.resource_record.record[new_resource.id] = new_resource
                            else:  # abort, backlog full
                                print("Backlog is full.")
                                # exit(1)

                        self.extra_info.new_resource_comes()

            reward = -1

            # add new jobs
            self.seq_idx += 1

        elif status == 'Allocate':
            self.resource_record.record[self.resource_slot.slot[a].id] = self.resource_slot.slot[a]
            self.resource_slot.slot[a] = None

            # dequeue backlog
            if self.resource_backlog.curr_size > 0:
                self.resource_slot.slot[a] = self.resource_backlog.backlog[0]  # if backlog empty, it will be 0
                self.resource_backlog.backlog[: -1] = self.resource_backlog.backlog[1:]
                self.resource_backlog.backlog[-1] = None
                self.resource_backlog.curr_size -= 1

        ob = self.observe()

        info = self.resource_record

        if done:
            self.seq_idx = 0
            reward = 1000

            if not repeat:
                self.seq_no = (self.seq_no + 1) % self.pa.num_ex

            self.reset()

        if self.render:
            self.plot_state()

        return ob, reward, done, info

    def reset(self):
        self.seq_idx = 0
        self.curr_time = 0

        # initialize system
        self.resource = Resource(self.pa)
        self.resource_slot = DisasterResourceSlot(self.pa)
        self.resource_backlog = DisasterResourceBacklog(self.pa)
        self.resource_record = DisasterResourceRecord()
        self.extra_info = ExtraInfo(self.pa)




class DisasterResource:
    def __init__(self, res_vec, resource_len, resource_id, enter_time ,weight):
        self.id = resource_id
        self.res_vec = res_vec
        self.len = resource_len
        self.enter_time = enter_time
        self.start_time = -1  # not being allocated
        self.finish_time = -1
        self.weight = weight


class DisasterResourceSlot:
    def __init__(self, pa):
        self.slot = [None] * pa.num_nw


class DisasterResourceBacklog:
    def __init__(self, pa):
        self.backlog = [None] * pa.backlog_size
        self.curr_size = 0

class DisasterResourceRecord:
    def __init__(self):
        self.record = {}



class Resource:
    def __init__(self, pa):
        self.num_res = pa.num_res
        self.time_horizon = pa.time_horizon
        self.res_slot = pa.res_slot

        self.avbl_slot = np.ones((self.time_horizon, self.num_res)) * self.res_slot

        self.running_resource = []

        # colormap for graphical representation
        self.colormap = np.arange(1 / float(pa.job_num_cap), 1, 1 / float(pa.job_num_cap))
        np.random.shuffle(self.colormap)

        # graphical representation
        self.canvas = np.zeros((pa.num_res, pa.time_horizon, pa.res_slot))

    def allocate_resource(self, resource, curr_time):

        allocated = False

        for t in range(0, self.time_horizon - resource.len):

            new_avbl_res = self.avbl_slot[t: t + resource.len, :] - resource.res_vec

            if np.all(new_avbl_res[:] >= 0):

                allocated = True

                self.avbl_slot[t: t + resource.len, :] = new_avbl_res
                resource.start_time = curr_time + t
                resource.finish_time = resource.start_time + resource.len

                self.running_resource.append(resource)

                # update graphical representation

                used_color = np.unique(self.canvas[:])
                # WARNING: there should be enough colors in the color map
                for color in self.colormap:
                    if color not in used_color:
                        new_color = color
                        break

                assert resource.start_time != -1
                assert resource.finish_time != -1
                assert resource.finish_time > resource.start_time
                canvas_start_time = resource.start_time - curr_time
                canvas_end_time = resource.finish_time - curr_time

                for res in range(self.num_res):
                    for i in range(canvas_start_time, canvas_end_time):
                        avbl_slot = np.where(self.canvas[res, i, :] == 0)[0]
                        self.canvas[res, i, avbl_slot[: resource.res_vec[res]]] = new_color

                break

        return allocated

    def time_proceed(self, curr_time):

        self.avbl_slot[:-1, :] = self.avbl_slot[1:, :]
        self.avbl_slot[-1, :] = self.res_slot

        for resource in self.running_resource:

            if resource.finish_time <= curr_time:
                self.running_resource.remove(resource)

        # update graphical representation

        self.canvas[:, :-1, :] = self.canvas[:, 1:, :]
        self.canvas[:, -1, :] = 0

# class ResourceAllocator:
#     def __init__(self, num_res, time_horizon, res_slot):
#         self.num_res = num_res
#         self.time_horizon = time_horizon
#         self.res_slot = res_slot
#
#         self.avbl_slot = np.ones((self.time_horizon, self.num_res)) * self.res_slot
#
#         self.allocated_resources = []
#
#     def allocate_resources(self, resource_demand, job_len, curr_time):
#
#         allocated = False
#
#         for t in range(0, self.time_horizon - job_len):
#
#             new_avbl_res = self.avbl_slot[t: t + job_len, :] - resource_demand
#
#             if np.all(new_avbl_res[:] >= 0):
#                 allocated = True
#
#                 self.avbl_slot[t: t + job_len, :] = new_avbl_res
#
#                 # Update the list of allocated resources
#                 self.allocated_resources.append({
#                     'start_time': curr_time + t,
#                     'finish_time': curr_time + t + job_len,
#                     'resource_demand': resource_demand
#                 })
#
#                 break
#
#         return allocated
#
#     def time_proceed(self, curr_time):
#
#         self.avbl_slot[:-1, :] = self.avbl_slot[1:, :]
#         self.avbl_slot[-1, :] = self.res_slot
#
#         self.allocated_resources = [job for job in self.allocated_resources if job['finish_time'] > curr_time]




class ExtraInfo:
    def __init__(self, pa):
        self.time_since_last_new_resource = 0
        self.max_tracking_time_since_last_resource = pa.max_track_since_new

    def new_resource_comes(self):
        self.time_since_last_new_resource = 0

    def time_proceed(self):
        if self.time_since_last_new_resource < self.max_tracking_time_since_last_resource:
            self.time_since_last_new_resource += 1

# ==========================================================================
# ------------------------------- Unit Tests -------------------------------
# ==========================================================================


def test_backlog():
    pa = parameters.Parameters()
    pa.num_nw = 5
    pa.simu_len = 50
    pa.num_ex = 10
    pa.new_job_rate = 1
    pa.compute_dependent_parameters()

    env = Env(pa, render=False, repre='image')

    env.step(5)
    env.step(5)
    env.step(5)
    env.step(5)
    env.step(5)

    env.step(5)
    assert env.job_backlog.backlog[0] is not None
    assert env.job_backlog.backlog[1] is None
    print("New job is backlogged.")

    env.step(5)
    env.step(5)
    env.step(5)
    env.step(5)

    job = env.job_backlog.backlog[0]
    env.step(0)
    assert env.job_slot.slot[0] == job

    job = env.job_backlog.backlog[0]
    env.step(0)
    assert env.job_slot.slot[0] == job

    job = env.job_backlog.backlog[0]
    env.step(1)
    assert env.job_slot.slot[1] == job

    job = env.job_backlog.backlog[0]
    env.step(1)
    assert env.job_slot.slot[1] == job

    env.step(5)

    job = env.job_backlog.backlog[0]
    env.step(3)
    assert env.job_slot.slot[3] == job

    print("- Backlog test passed -")


def test_compact_speed():

    pa = parameters.Parameters()
    pa.simu_len = 50
    pa.num_ex = 10
    pa.new_job_rate = 0.3
    pa.compute_dependent_parameters()

    env = Env(pa, render=False, repre='compact')

    import other_agents
    import time

    start_time = time.time()
    for i in range(100000):
        a = other_agents.get_genetic_algorithm_action(env.resource, env.resource_slot)
        env.step(a)
    end_time = time.time()
    print("- Elapsed time: ", end_time - start_time, "sec -")


def test_image_speed():

    pa = parameters.Parameters()
    pa.simu_len = 50
    pa.num_ex = 10
    pa.new_job_rate = 0.3
    pa.compute_dependent_parameters()

    env = Env(pa, render=False, repre='image')

    import other_agents
    import time

    start_time = time.time()
    for i in range(100000):
        a = other_agents.get_genetic_algorithm_packer_action(env.resource, env.resource_slot)
        env.step(a)
    end_time = time.time()
    print("- Elapsed time: ", end_time - start_time, "sec -")


if __name__ == '__main__':
    test_backlog()
    test_compact_speed()
    test_image_speed()
